#!/usr/bin/env python3
# -*- coding: utf-8 -*-



import sys, cv2
print("== DEBUG from rtsp.py ==")
print("Python:", sys.version)
print("sys.executable:", sys.executable)
print("cv2:", cv2.__file__)

from cv_bridge import CvBridge, CvBridgeError
print("cv_bridge OK in rtsp.py:", CvBridge)

import rospy
from sensor_msgs.msg import Image
import time

# ---------- 屏幕尺寸 ----------
SCREEN_W = 1920
SCREEN_H = 1200

# ---------- RTSP 取流 ----------
RTSP_URL = "rtsp://192.168.1.50:8554/air"
pipeline = (
    f"rtspsrc location={RTSP_URL} protocols=tcp latency=0 ! "
    "rtph264depay ! h264parse ! avdec_h264 ! "
    "videoconvert ! appsink drop=true max-buffers=1 sync=false"
)

bridge = CvBridge()
latest_ros_img = None   # 最新 ROS 图像（主画面）

def ros_image_callback(msg):
    global latest_ros_img
    try:
        img = bridge.imgmsg_to_cv2(msg, "bgr8")
        latest_ros_img = img
    except CvBridgeError as e:
        rospy.logerr("CvBridge Error: %s", e)

def open_rtsp():
    """尝试打开 RTSP，成功返回 cap，失败返回 None"""
    cap = cv2.VideoCapture(pipeline, cv2.CAP_GSTREAMER)
    if not cap.isOpened():
        rospy.logwarn("无法打开 RTSP/GStreamer 管道，将在 1 秒后重试")
        cap.release()
        return None
    rospy.loginfo("RTSP 流打开成功")
    return cap

def main():
    global latest_ros_img

    rospy.init_node("ros_rtsp_fusion_viewer", anonymous=True)

    # 订阅 ROS 主画面
    rospy.Subscriber("/shouchi/image", Image, ros_image_callback, queue_size=1)

    win = "ROS + RTSP"
    cv2.namedWindow(win, cv2.WINDOW_NORMAL)

    cap = None
    last_retry_time = 0.0
    moved = False  # 是否已经把窗口摆到屏幕右上角

    rate = rospy.Rate(30)  # 30Hz 刷新

    while not rospy.is_shutdown():
        # 没有主画面先等着
        if latest_ros_img is None:
            rate.sleep()
            continue

        # 拿一帧主画面
        frame = latest_ros_img.copy()
        h, w = frame.shape[:2]

        # ========== 处理 RTSP 副画面 ==========
        now = time.time()

        # RTSP 没打开或挂了 -> 每秒重试一次
        if cap is None or not cap.isOpened():
            if now - last_retry_time >= 1.0:
                cap = open_rtsp()
                last_retry_time = now

        pip_frame = None

        if cap is not None and cap.isOpened():
            ok, rtsp_frame = cap.read()
            if not ok or rtsp_frame is None:
                rospy.logwarn("RTSP 帧读取失败，将重新连接")
                cap.release()
                cap = None
                last_retry_time = now
            else:
                # 按主画面高度的 1/4 做副画面高度
                pip_h = int(h * 0.25)
                rh, rw = rtsp_frame.shape[:2]
                pip_w = int(float(rw) / rh * pip_h)

                pip_frame = cv2.resize(
                    rtsp_frame, (pip_w, pip_h), interpolation=cv2.INTER_AREA
                )

        # ========== 把 RTSP 副画面贴到主画面右上角 ==========
        if pip_frame is not None:
            ph, pw = pip_frame.shape[:2]
            margin = 10
            x1 = w - pw - margin   # 右上角 x
            y1 = margin            # 右上角 y

            # 边界保护（如果图像尺寸比较怪）
            if x1 < 0:
                x1 = 0
            if y1 < 0:
                y1 = 0
            if x1 + pw > w:
                pw = w - x1
                pip_frame = pip_frame[:, :pw]
            if y1 + ph > h:
                ph = h - y1
                pip_frame = pip_frame[:ph, :]

            frame[y1:y1 + ph, x1:x1 + pw] = pip_frame

        # ========== 处理窗口位置 / 大小 ==========
        if not moved:
            # 根据屏幕分辨率算一个缩放比例，让窗口别超出屏幕
            scale_w = float(SCREEN_W) / w
            scale_h = float(SCREEN_H) / h
            scale = min(1.0, scale_w, scale_h)  # 不放大，只缩小

            win_w = int(w * scale)
            win_h = int(h * scale)

            cv2.resizeWindow(win, win_w, win_h)
            cv2.moveWindow(win, SCREEN_W - win_w, 0)  # 窗口放到屏幕右上角

            moved = True

        # ========== 显示最终画面 ==========
        cv2.imshow(win, frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        rate.sleep()

    # 收尾
    if cap is not None:
        cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
